from datetime import datetime
from typing import List, Any
from ..core.mcp import MCPClient
from ..core.common import AgentConfig, ModelParameter, MCPServerConfig
from ..core.schemas import  TextBlock, Message, ResponseMessage
from ..agenthub import DummyAgent
from ..schemas import ChatRequest,ChatErrorData

from scorpio.core.common import get_logger
logger = get_logger(__name__)


chat_history=[]
chat_step_code={
    "initialize": 1001,
    "ping": 1002,
    "create_model": 1005,
}
def _default_chat_error_response(step: str, data: Any)->ChatErrorData:
    return ChatErrorData(
        code=chat_step_code[step],
        message=step,
        data=data
    )

def convert_response_messages(messages: List[Message])-> List[ResponseMessage]:
    result: List[ResponseMessage] = []
    for idx, message in enumerate(messages):
        if message.role == "assistant":
            result.append(message.to_response_message(idx))
        elif  message.role == "tool":
            result.append(message.to_response_message(idx))
        else:
            pass
    return result

async def chat_with_agent(request: ChatRequest, mcp_config: MCPServerConfig, chosen_model: ModelParameter)->str|ChatErrorData:
    agent_config=AgentConfig(mcp_config=mcp_config, model_parameter=chosen_model)
    shared = {
        "question": request.question, 
        "system_prompt": "simple.md",
        "context": {
            "agent":{
                "name": "dummy",
                "config": agent_config,
                "mcp_server": request.mcp_server,
            },
            "messages": [],
            "events": [],
        }
    }
    agent=DummyAgent(shared=shared)
    agent.create_flow()
    step=await agent.run()
    logger.info(f"PocketFlow end with step: {step}")
    messages: List[Message]=shared["context"].get("messages")
    result=convert_response_messages(messages)
    return result
        


# 移除原硬编码的models_cfg，改为通过参数接收
async def chat_with_mcp(request: ChatRequest, mcp_config: MCPServerConfig, chosen_model: ModelParameter)->str|ChatErrorData:
    mcp_server=request.mcp_server
    logger.info(f"Server {mcp_server}: transport protocol {mcp_config.transportType}...")

    cli = MCPClient(
        mcp_server,
        mcp_config,
        model_param=chosen_model )
    # print("Get MCP Client",chosen_model)
    # start: initialize MCP server connection session
    ok=await cli.start()
    if not ok:
        return _default_chat_error_response("initialize", f" MCP server {mcp_server}: failed to initialize connection session")
    
    # ping: check MCP server connection state
    ok=await cli.ping()
    if not ok:
        return _default_chat_error_response("ping", f"MCP server {mcp_server}: check connection state failed")
    try:
        content=TextBlock()
        content.text="You are a assistant with ability to use tool with function call and summary result for question."
        system_message=Message(role="system", content=[content])
        messages: List[Message]=[system_message]

        response: List[Message]=await cli.process(messages, request.question)
        result=convert_response_messages(response)
        # 新增：保存聊天记录到内存列表
        chat_history.append({
            "id": len(chat_history) + 1,
            "question": request.question,
            "answer": result,
            "timestamp": datetime.now()
        })
    finally:
        await cli.cleanup()
    return result

